# 导入必要的模块
from modelscope import AutoModelForCausalLM, AutoTokenizer

# 设定设备为GPU（若可用）
# 注意：确保你的环境支持CUDA，并且PyTorch已正确配置以使用GPU
device = "cuda" # 设定模型加载的设备为GPU

# 从预训练模型路径加载因果语言模型（Causal Language Model）
# "device_map='auto'" 表示模型会被自动分配到可用的GPU上，如果存在多个GPU，则模型会被分散加载
model = AutoModelForCausalLM.from_pretrained(
    "G:/QwenModel/qwen/Qwen1___5-1___8B-Chat", # 预训练模型的本地路径
    device_map="auto" # 自动分配模型到GPU
)

# 从相同的预训练模型路径加载分词器（Tokenizer）
tokenizer = AutoTokenizer.from_pretrained("G:/QwenModel/qwen/Qwen1___5-1___8B-Chat")

# 定义用户询问的内容
prompt = input("请输入你想问的问题：")
# 构建对话消息列表，包括系统角色和用户角色的消息
# 系统角色消息设定模型为“有帮助的助手”
# 用户角色消息为用户的询问内容
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": prompt}
]

# 使用分词器应用聊天模板，将消息转换成模型可以理解的格式
# "tokenize=False" 表示不进行分词化，"add_generation_prompt=True" 添加生成提示
text = tokenizer.apply_chat_template(
    messages,
    tokenize= False, # 不进行分词化
    add_generation_prompt=True # 添加生成提示
)

# 使用分词器对文本进行编码，转换成模型输入所需的张量格式
# "return_tensors='pt'" 表示返回PyTorch张量
model_inputs = tokenizer([text], return_tensors="pt")

# 将模型输入数据移动到预先设定的设备（GPU）
model_inputs = model_inputs.to(device)

# 使用模型进行生成，"generate()" 方法用于文本生成任务
# "max_new_tokens=512" 表示最大新生成的令牌数量为512
generated_ids = model.generate(
    model_inputs.input_ids, # 输入ID
    max_new_tokens=512 # 最大新生成的令牌数量
)

# 分离生成的ID与原始输入ID，只保留新生成的部分
# 这里使用列表推导式遍历输入和输出ID，计算新生成的ID部分
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

# 使用分词器解码生成的ID，转换回人类可读的文本
# "skip_special_tokens=True" 表示在解码时跳过特殊令牌
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)
# `response` 现在包含了由模型生成的回答